package org.fitweb.fixture.timing;

import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

public class TimedExecutor<T> {

    private final int timeout;
    private final TimeUnit unit;

    private TimedExecutor(int timeout, TimeUnit unit) {
        this.timeout = timeout;
        this.unit = unit;
    }

    public static <X> TimedExecutor<X> waitFor(int timeout, TimeUnit unit) {
        return new TimedExecutor<X>(timeout, unit);
    }

    public T on(Callable<T> callable) throws TimeoutException {
        ExecutorService executorService = Executors.newSingleThreadExecutor();
        Future<T> future = executorService.submit(callable);
        try {
            return future.get(timeout, unit);
        } catch(ExecutionException e) {
            throw new RuntimeException("TimedExecutor task threw exception", e);
        } catch(InterruptedException e) {
            throw new RuntimeException("TimedExecutor task interrupted", e);
        } finally {
            future.cancel(true);
            executorService.shutdownNow();
        }
    }
}